import cv2
import numpy as np
import matplotlib.pyplot as plt


def cal_distance(pa, pb):
    """计算距离"""
    from math import sqrt
    dis = sqrt((pa[0] - pb[0]) ** 2 + (pa[1] - pb[1]) ** 2)
    return dis

# d0为半径代表从什么时候这个权重矩阵开始从0到1
def get_gauss_lowpass(shape, d0):
    """返回高斯低通滤波矩阵"""
    mat = np.zeros(shape)
    center_point = (shape[0] // 2, shape[1] // 2)
    for i in range(shape[0]):
        for j in range(shape[1]):
            dis = cal_distance(center_point, (i, j))
            mat[i, j] = np.exp(-(dis ** 2) / (2 * (d0 ** 2)))

    return mat

# d0为半径代表从什么时候这个权重矩阵开始从1到0
def get_gauss_highpass(shape, d0):
    """返回高斯高通滤波矩阵"""
    return 1 - get_gauss_lowpass(shape, d0)


def test():
    image = cv2.imread('img/girl.png', cv2.IMREAD_GRAYSCALE)
    # 变换到频率域
    # fft2代表二位傅里叶变换 得到没有将0频放到中间的频谱
    f_image = np.fft.fft2(image)
    # fftshift 就把0频给平移到中间了
    f_image = np.fft.fftshift(f_image)
    # 由于中间的低频能量太高了，显示出来就是中间一个白点周围全是黑的
    # 于是我们在这里取一个对数将低频的能量进行一个削弱就能只管的看到效果了
    spectrum = 20 * np.log(np.abs(f_image))
    # 频率域滤波
    # filter_mat = get_gauss_lowpass(image.shape,50)
    filter_mat = get_gauss_highpass(image.shape,10)
    f_filter_image = f_image * filter_mat
    filter_spectrum = 20 * np.log(np.abs(f_filter_image))
    # 反变换回空间域
    # 反向移动之前的0频移动
    f_filter_image = np.fft.ifftshift(f_filter_image)
    # 反变换为图像
    filter_image = np.fft.ifft2(f_filter_image)
    filter_image = np.abs(filter_image)
    # 显示结果图像
    plt.subplot(221), plt.imshow(image, cmap="gray")
    plt.title("Src image"), plt.xticks([]), plt.yticks([])
    plt.subplot(222), plt.imshow(spectrum, cmap="bwr")
    plt.title('Src spectrum'), plt.xticks([]), plt.yticks([])
    plt.subplot(223), plt.imshow(filter_image,cmap='gray')
    plt.title("Filter image"), plt.xticks([]), plt.yticks([])
    plt.subplot(224), plt.imshow(filter_spectrum, cmap="bwr")
    plt.title('Filter spectrum'), plt.xticks([]), plt.yticks([])
    plt.show()


test()